import os
import cv2
import numpy as np
import matplotlib.pyplot as plt
import paddle
from paddle.io import Dataset
import parameter

img_dim_w = parameter.img_dim_w
img_dim_h = parameter.img_dim_h


"""
若报错，尝试安装下面版本的包

pip install numpy==1.15.4
pip install matplotlib==2.2.2 --user
"""

#准备数据，定义Reader()
PATH = parameter.PATH
class DataGenerater(Dataset):
    """
    数据集定义
    """
    def __init__(self,path=PATH):
        """
        构造函数
        """
        super(DataGenerater, self).__init__()
        self.dir = path
        self.datalist = os.listdir(PATH)
        self.image_size = (img_dim_w,img_dim_w)
    
    # 每次迭代时返回数据和对应的标签
    def __getitem__(self, idx):
        while True:
            try:
                img = self._load_img(self.dir + self.datalist[idx])
                break
            except:
                idx += 1
        return img
            
    # 返回整个数据集的总数
    def __len__(self):
        return len(self.datalist)
    
    def _load_img(self, path):
        """
        统一的图像处理接口封装，用于规整图像大小和通道
        """
        img = cv2.imread(path)
        img = cv2.resize(img,self.image_size)
        img = img.transpose()
        img = img.astype('float32')
        img = img / 255.
        
        return img



if __name__ == "__main__":
    train_dataset = DataGenerater()
    train_loader = paddle.io.DataLoader(train_dataset,
                        batch_size=100,
                        shuffle=True,
                        drop_last=True)
    
    for batch_id, data in enumerate(train_loader()):
        plt.figure(figsize=(15,15))
        try:
            for i in range(100):
                image = np.array(data[i].transpose((2,1,0)))
                plt.subplot(10, 10, i + 1)
                plt.imshow(image, vmin=-1, vmax=1)
                plt.axis('off')
                plt.xticks([])
                plt.yticks([])
                plt.subplots_adjust(wspace=0.1, hspace=0.1)
            plt.suptitle('\n Training Images',fontsize=30)
            plt.show()
            break
        except IOError:
            print(IOError)